"""
@Filename       : preprocess.py
@Create Time    : 2021/11/23 11:06
@Author         : Rylynn
@Description    : 

"""
import os
import pickle as pkl
import numpy as np
import networkx as nx

def process_blog_catalog(root_path):
    node_list = []
    g = nx.DiGraph()
    with open(os.path.join(root_path, 'BlogCatalog3/raw/nodes.csv')) as f:
        for line in f.readlines():
            node = eval(line)
            g.add_node(node)

    with open(os.path.join(root_path, 'BlogCatalog3/raw/edges.csv')) as f:
        for line in f.readlines():
            source, target = line.strip().split(',')
            source = eval(source)
            target = eval(target)
            g.add_edge(source, target)

    label_dict = {}
    with open(os.path.join(root_path, 'BlogCatalog3/raw/group-edges.csv')) as f:
        for line in f.readlines():
            node, label = line.strip().split(',')
            node = eval(node)
            label = eval(label)
            label_dict[node] = label

    label_tuple = label_dict.items()
    label_tuple = sorted(label_tuple, key=lambda x: x[0])
    labels = list(map(lambda x: x[0], label_tuple))
    assert len(labels) == g.number_of_nodes(), "The number of nodes unmatches the number of labels"

    A = nx.to_scipy_sparse_matrix(g)
    pkl.dump({"adj": A, "label": label}, open(os.path.join(root_path, 'BlogCatalog3/preprocessed/graph.pkl'), 'wb'))




process_blog_catalog('../../data')